#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RocketMQ 批量生产者示例
演示如何使用批量方式发送消息
"""

import time
import json
import logging
from rocketmq.client import Producer, Message
from config import ROCKETMQ_CONFIG, MESSAGE_CONFIG, LOGGING_CONFIG

# 配置日志
logging.basicConfig(
    level=getattr(logging, LOGGING_CONFIG["level"]), 
    format=LOGGING_CONFIG["format"]
)
logger = logging.getLogger(__name__)


class BatchRocketMQProducer:
    """批量 RocketMQ 生产者类"""
    
    def __init__(self):
        self.producer = None
        self.namesrv_addr = ROCKETMQ_CONFIG["namesrv_addr"]
        self.access_key = ROCKETMQ_CONFIG["access_key"]
        self.secret_key = ROCKETMQ_CONFIG["secret_key"]
        self.group_name = ROCKETMQ_CONFIG["producer_group"]
        self.topic = ROCKETMQ_CONFIG["default_topic"]
        self.batch_size = 10  # 默认批量大小
    
    def start(self):
        """启动生产者"""
        try:
            logger.info("=== 启动批量 RocketMQ 生产者 ===")
            logger.info(f"NameServer: {self.namesrv_addr}")
            logger.info(f"生产者组: {self.group_name}")
            logger.info(f"主题: {self.topic}")
            logger.info(f"批量大小: {self.batch_size}")
            
            # 创建生产者
            self.producer = Producer(self.group_name)
            self.producer.set_name_server_address(self.namesrv_addr)
            self.producer.set_session_credentials(self.access_key, self.secret_key, "ALIYUN")
            
            # 启动生产者
            self.producer.start()
            logger.info("✓ 批量生产者启动成功")
            return True
            
        except Exception as e:
            logger.error(f"✗ 批量生产者启动失败: {e}")
            return False
    
    def create_message(self, message_data, tags=None, keys=None):
        """创建消息对象"""
        msg = Message(self.topic)
        
        # 设置消息内容
        if isinstance(message_data, dict):
            body = json.dumps(message_data, ensure_ascii=False)
        else:
            body = str(message_data)
        
        msg.set_body(body.encode(MESSAGE_CONFIG["encoding"]))
        
        # 设置消息标签
        if tags:
            msg.set_tags(tags)
        else:
            msg.set_tags(MESSAGE_CONFIG["default_tags"])
        
        # 设置消息键
        if keys:
            msg.set_keys(keys)
        
        return msg
    
    def send_batch(self, message_list, batch_id):
        """发送一批消息"""
        try:
            if not message_list:
                return 0
            
            logger.info(f"发送批次 {batch_id}: {len(message_list)} 条消息")
            
            # 创建批量消息列表
            batch_messages = []
            for i, msg_data in enumerate(message_list):
                msg = self.create_message(
                    message_data=msg_data,
                    tags=f"batch_{batch_id}_{i + 1}",
                    keys=f"batch_key_{batch_id}_{i + 1}"
                )
                batch_messages.append(msg)
            
            # 发送批量消息
            start_time = time.time()
            send_result = self.producer.send_batch(batch_messages)
            end_time = time.time()
            
            # 计算发送时间
            send_time = (end_time - start_time) * 1000  # 转换为毫秒
            
            # 打印发送结果
            logger.info(f"✓ 批次 {batch_id} 发送成功 (耗时: {send_time:.2f}ms)")
            logger.info(f"  - 消息数量: {len(message_list)}")
            logger.info(f"  - 平均每条消息: {send_time / len(message_list):.2f}ms")
            
            # 兼容性处理
            if hasattr(send_result, 'msg_id'):
                logger.info(f"  - 批量MsgId: {send_result.msg_id}")
            if hasattr(send_result, 'offset'):
                logger.info(f"  - 偏移量: {send_result.offset}")
            if hasattr(send_result, 'status'):
                logger.info(f"  - 状态: {send_result.status}")
            
            return len(message_list)
            
        except Exception as e:
            logger.error(f"✗ 批次 {batch_id} 发送失败: {e}")
            return 0
    
    def send_large_batch(self, total_messages, batch_size=None):
        """发送大批量消息"""
        if batch_size is None:
            batch_size = self.batch_size
        
        logger.info(f"=== 开始批量发送 {total_messages} 条消息 (批次大小: {batch_size}) ===")
        
        # 创建所有消息
        all_messages = self.create_batch_demo_messages(total_messages)
        
        # 分批发送
        total_sent = 0
        total_batches = (len(all_messages) + batch_size - 1) // batch_size
        
        start_time = time.time()
        
        for batch_idx in range(total_batches):
            start_idx = batch_idx * batch_size
            end_idx = min(start_idx + batch_size, len(all_messages))
            batch_messages = all_messages[start_idx:end_idx]
            
            # 发送当前批次
            sent_count = self.send_batch(batch_messages, batch_idx + 1)
            total_sent += sent_count
            
            # 批次间隔
            if batch_idx < total_batches - 1:
                time.sleep(0.5)  # 500ms间隔
        
        end_time = time.time()
        total_time = end_time - start_time
        
        # 统计信息
        logger.info(f"=== 批量发送完成 ===")
        logger.info(f"总消息数: {total_messages}")
        logger.info(f"成功发送: {total_sent}")
        logger.info(f"总耗时: {total_time:.2f}秒")
        logger.info(f"平均速度: {total_sent / total_time:.1f} 消息/秒")
        logger.info(f"平均每批次: {total_time / total_batches:.3f}秒")
        
        return total_sent
    
    def create_batch_demo_messages(self, count):
        """创建批量演示消息"""
        messages = []
        
        # 不同类型的消息模板
        message_templates = [
            {"type": "batch_text", "content": "Hello Batch RocketMQ!"},
            {"type": "batch_json", "data": {"user": "batch_demo", "action": "send_batch"}},
            {"type": "batch_number", "value": 999},
            {"type": "batch_list", "items": [100, 200, 300, 400, 500]},
            {"type": "batch_boolean", "flag": True},
            {"type": "batch_object", "obj": {"name": "batch", "version": "1.0"}},
        ]
        
        for i in range(count):
            template = message_templates[i % len(message_templates)]
            msg_data = {
                "id": i + 1,
                "timestamp": int(time.time()),
                "message": f"批量演示消息 #{i + 1}",
                "data": template,
                "source": "python_batch_demo",
                "send_type": "batch",
                "batch_index": i + 1
            }
            messages.append(msg_data)
        
        return messages
    
    def shutdown(self):
        """关闭生产者"""
        if self.producer:
            self.producer.shutdown()
            logger.info("批量生产者已关闭")


def main():
    """主函数"""
    logger.info("开始 RocketMQ 批量生产者演示...")
    
    # 创建批量生产者
    producer = BatchRocketMQProducer()
    
    try:
        # 启动生产者
        if not producer.start():
            return
        
        # 发送不同大小的批量消息
        batch_sizes = [5, 10, 20, 50]
        total_messages = 100
        
        for batch_size in batch_sizes:
            logger.info(f"\n{'='*50}")
            logger.info(f"测试批量大小: {batch_size}")
            logger.info(f"{'='*50}")
            
            # 发送批量消息
            success_count = producer.send_large_batch(total_messages, batch_size)
            
            # 等待一下再进行下一轮测试
            if batch_size != batch_sizes[-1]:  # 不是最后一轮
                logger.info("等待 3 秒后进行下一轮测试...")
                time.sleep(3)
        
        # 总结
        logger.info(f"\n{'='*50}")
        logger.info("批量生产者演示完成")
        logger.info(f"{'='*50}")
        logger.info("✓ 批量生产者演示成功完成")
        logger.info("提示: 批量发送可以提高消息发送效率")
        
    except KeyboardInterrupt:
        logger.info("用户中断演示")
    except Exception as e:
        logger.error(f"演示过程中发生错误: {e}")
    finally:
        producer.shutdown()


if __name__ == "__main__":
    main() 